from langchain.chains import LLMChain, SequentialChain, TransformChain
from langchain.prompts import ChatPromptTemplate
from langchain.schema import SystemMessage  # 请确认使用正确的消息类
from langchain_openai import ChatOpenAI
from langchain.tools import tool
import langchain
import subprocess
import matplotlib.pyplot as plt
import numpy as np
import re
import json
import os
from dotenv import load_dotenv
from spice_gen import generate_bandpass_filter

langchain.verbose = True
langchain.debug = True

load_dotenv()

# @tool
def generate_bandpass_filter_tool(
    order: int,
    highpass_r: float,
    highpass_c: float,
    lowpass_r: float,
    lowpass_c: float,
    filename: str = "BPF.sp"
) -> str:
    """
    生成一个带通滤波器（Band-Pass Filter, BPF）的网表文件，用于模拟高通 + 低通滤波器组合电路。

    参数说明：
    - order (int): 滤波器的阶数，影响滤波器的陡峭程度（例如 2 表示二阶滤波器）。
    - highpass_r (float): 高通部分的电阻值（单位：欧姆）。
    - highpass_c (float): 高通部分的电容值（单位：法拉）。
    - lowpass_r (float): 低通部分的电阻值（单位：欧姆）。
    - lowpass_c (float): 低通部分的电容值（单位：法拉）。
    - filename (str): 输出的 SPICE 网表文件名，默认是 "BPF.sp"。

    返回值：
    - str: 成功生成网表文件的提示信息，包括文件名。
    """
    # 调用实际生成滤波器网表的函数
    generate_bandpass_filter(
        order=order,
        highpass_r=highpass_r,
        highpass_c=highpass_c,
        lowpass_r=lowpass_r,
        lowpass_c=lowpass_c,
        filename=filename
    )
    return f"带通滤波器已成功生成，保存为文件: {filename}"

# @tool
def run_simulation(text: str) -> str:
    """
    执行HSPICE仿真任务

    输入要求：一个空字符串

    成功返回：
    "仿真完成：output.lis"

    失败返回：
    "仿真失败：[错误信息]"

    输出文件：
    - output.lis : 文本报告
    - output.mt0  : 频响数据
    """
    print("开始仿真...")
    try:
        subprocess.run(
            ["hspice", "-i", "BPF.sp", "-o", "output"],
            check=True,
            capture_output=True
        )
        return "仿真完成：output.lis"
    except subprocess.CalledProcessError as e:
        return f"仿真失败：{e.stderr.decode()}"

# @tool
def analyze_results(text: str) -> str:
    """
    解析仿真结果并生成报告

    输入要求：一个空字符串

    成功返回：
    JSON格式：
    {
        "BW": 实际带宽（单位Hz）,
        "F0": 中心频率（单位Hz）,
        "f_low": 低截止频率（单位Hz）,
        "f_high": 高截止频率（单位Hz）,
        "status": "指标验证结果（达标/不达标）"
    }

    失败返回：
    {"error": "[错误描述]"}
    """
    print("开始结果分析...")
    def parse_value(val_str):
        """将带单位的字符串转换为浮点数"""
        unit_multipliers = {'k': 1e3, 'K': 1e3, 'M': 1e6, 'G': 1e9}
        match = re.match(r'([\d\.Ee+-]+)\s*([kKmMgG]?)', val_str.strip())
        if not match:
            raise ValueError(f"无法解析数值：{val_str}")
        number, unit = match.groups()
        return float(number) * unit_multipliers.get(unit, 1)

    try:
        # 解析测量结果
        with open("output.lis") as f:
            content = f.read()

        # 修改正则表达式兼容科学记数法（E/e）
        measures = {
            'BW': parse_value(re.search(r'bw\s*=\s*([\d\.Ee+-]+\s*[kKmMgG]?)', content, re.IGNORECASE).group(1)),
            'F0': parse_value(re.search(r'f_center\s*=\s*([\d\.Ee+-]+\s*[kKmMgG]?)', content, re.IGNORECASE).group(1)),
            'f_low': parse_value(re.search(r'f_low\s*=\s*([\d\.Ee+-]+\s*[kKmMgG]?)', content, re.IGNORECASE).group(1)),
            'f_high': parse_value(re.search(r'f_high\s*=\s*([\d\.Ee+-]+\s*[kKmMgG]?)', content, re.IGNORECASE).group(1))
        }
        print(f"##########\n测量结果：BW={measures['BW']} Hz, F0={measures['F0']} Hz\n, f_low={measures['f_low']} Hz, f_high={measures['f_high']} Hz\n##########")
        
        # 读取设计要求
        with open("design_spec.json") as f:
            spec = json.load(f)

        # 允许的误差范围
        bw_tolerance = 0.1 * spec["bandwidth"]
        f0_tolerance = 0.1 * spec["center_frequency"]

        status = ("达标" if 
            (abs(measures['BW'] - spec["bandwidth"])) < bw_tolerance and
            (abs(measures['F0'] - spec["center_frequency"]))< f0_tolerance
            else "不达标"
        )

        return json.dumps({
            "BW": measures['BW'],
            "F0": measures['F0'],
            "f_low": measures['f_low'],
            "f_high": measures['f_high'],
            "status": status
        }, ensure_ascii=False)
    except Exception as e:
        return json.dumps({"error": f"结果分析失败：{str(e)}"})

# def visualize(text: str) -> str:
#     print("开始可视化...")

#     try:
#         with open("output.ac0", "r", encoding="utf-8", errors="ignore") as f:
#             lines = f.readlines()

#         # Step 1: 找到包含 "HERTZ" 的行，作为 header 起始
#         header_line_idx = None
#         for i, line in enumerate(lines):
#             if "HERTZ" in line.upper():
#                 header_line_idx = i
#                 break

#         if header_line_idx is None:
#             return "可视化失败：找不到包含 HERTZ 的变量头部"

#         # Step 2: 提取接下来几行的变量名，横向读取
#         variable_lines = lines[header_line_idx + 1 : header_line_idx + 5]
#         variable_names = []
#         for line in variable_lines:
#             tokens = line.strip().lower().split()
#             variable_names.extend(tokens)
#         if "vdb(out)" not in variable_names:
#             return f"可视化失败：未找到 vdb(out) 列，提取到的列有：{variable_names}"

#         vdb_index = variable_names.index("vdb(out)")

#         # Step 3: 寻找数据起始行（跳过空行/非数值行）
#         data_start_idx = header_line_idx + 5
#         while data_start_idx < len(lines):
#             line = lines[data_start_idx]
#             if re.match(r"^\s*[\d\.\-Ee+]+", line):  # 有效数值行
#                 break
#             data_start_idx += 1

#         # Step 4: 读取数据
#         data = []
#         for line in lines[data_start_idx:]:
#             try:
#                 row = [float(x) for x in line.strip().split()]
#                 data.append(row)
#             except:
#                 continue

#         data = np.array(data)
#         if data.shape[1] <= vdb_index:
#             return f"可视化失败：数据列数不足，vdb(out) 列索引为 {vdb_index}"

#         freq = data[:, 0]
#         gain = data[:, vdb_index]

#         # Step 5: 画图
#         plt.figure(figsize=(10, 6))
#         plt.semilogx(freq, gain)
#         plt.title('Frequency Response: VDB(out)')
#         plt.xlabel('Frequency (Hz)')
#         plt.ylabel('Gain (dB)')
#         plt.grid(which='both', linestyle='--')
#         plt.tight_layout()
#         plt.savefig('response.png')
#         return "图表保存成功：response.png"

#     except Exception as e:
#         return f"可视化失败：{str(e)}"
def parse_json_response(text: str) -> str:
    """
    如果输入的 JSON 字符串被markdown代码块包裹（例如 ```json ... ```），则去除包裹部分返回纯JSON文本。
    """
    text = text.strip()
    if text.startswith("```json"):
        # 去掉起始的 ```json 行
        text = text[text.find("\n") + 1:]
    if text.endswith("```"):
        # 去掉末尾的 ```
        text = text[:text.rfind("```")]
    return text.strip()

def save_and_parse(text: str, file_name: str) -> dict:
    raw = text
    parsed = json.loads(parse_json_response(raw))
    with open(f"{file_name}.json", "w", encoding="utf-8") as f:
        json.dump(parsed, f, indent=2, ensure_ascii=False)
    return text



# --------------------- Chain定义 ---------------------

llm = ChatOpenAI(model="qwq-32b", temperature=0, streaming=True, verbose=True)

# 1. 需求解析Chain：将用户需求转换为设计规范
# 设计规范解析 + 保存 JSON 文件
design_spec_chain = SequentialChain(
    chains=[
        LLMChain(
            llm=llm,
            prompt=ChatPromptTemplate.from_messages([
                SystemMessage(content="""将用户需求转换为设计规范，输出严格遵循以下JSON格式：
                {
                    "center_frequency": 中心频率（单位Hz）,
                    "bandwidth": 带宽（单位Hz）,
                    "tolerance": 允许误差百分比（默认10%）
                }"""),
                ("human", "设计需求：{input}")
            ]),
            output_key="design_spec_raw"
        ),
        TransformChain(
            input_variables=["design_spec_raw"],
            output_variables=["design_spec"],
            transform=lambda inputs: {"design_spec": save_and_parse(inputs["design_spec_raw"], "design_spec")}
        )
    ],
    input_variables=["input"],
    output_variables=["design_spec"]
)


# 2. 参数计算Chain：根据设计规范确定滤波器参数，要求输出JSON中至少包含以下字段：
# - order: 滤波器阶数
# - highpass_r, highpass_c: 高通部分元件参数
# - lowpass_r, lowpass_c: 低通部分元件参数
# - lib: 工艺库路径
# - rationale: 参数计算和设计思路说明
calc_chain = LLMChain(
    llm=llm,
    prompt=ChatPromptTemplate.from_messages([
        SystemMessage(content="""你是一个资深微电子工程师，根据输入的设计要求计算无源RC带通滤波器的参数。
你能够很好地设计滤波器，这是因为你知道中心频率和带宽与滤波器的阶数、元件参数之间的关系。
人们通常使用简化的截止频率计算公式来估算RC滤波器的截止频率，但是由于RC滤波器的负载效应、无源低阶滤波器的滚降系数（Roll-off Factor）不理想等因素，会导致实际截止频率偏离理论值。特别是会导致带宽偏大。
你要意识到这些问题，并试图在迭代中解决他们。
当增大无源滤波器的阶数时，滚降系数将变得更加理想，带宽也会更接近理论值，但是这会引起更严重的负载效应，使得中心频率发生一定的偏移，这时可以通过改变RC参数来补偿。
实际应用中不会使用过高阶数的滤波器，但是这里不关注设计出来的滤波器的实际可用性，只要满足设计要求即可。
注意：
1. 带通滤波器是由高通滤波器和低通滤波器级联构成，高通和低通的拓扑结构有所不同。
2. 请考虑级联带来的负载效应，为了降低低通级带来的负载效应，可以将低通级的电阻值设置为远大于高通级的电阻值。
3. 你可以使用非常用的电阻和电容值来设计滤波器，任意数值都行。
4. 可能会遇到带宽过大的问题，这是因为低阶无源滤波器的滚降系数不理想，导致带宽偏大，可以增大阶数来解决这个问题。
5. 带宽可能难以满足要求，请至少满足中心频率的要求，仔细回忆中心频率的公式，并把RC参数带入此公式计算。
6. 请先根据要求计算出带通滤波器的两个截止频率，然后根据截止频率计算出高通和低通的RC参数。
7. 输出必须严格符合以下JSON结构，不要添加多余信息：
{
    "order": 滤波器阶数,
    "highpass_r": 高通电阻值（单位欧姆）,
    "highpass_c": 高通电容值（单位法拉）,
    "lowpass_r": 低通电阻值（单位欧姆）,
    "lowpass_c": 低通电容值（单位法拉）,
    "rationale": 参数计算原理说明
}
"""),
        ("human", "设计要求：{design_spec}")
    ]),
    output_key="component_params"
)

# 3. 仿真验证和结果处理Chain：包括网表生成、HSPICE仿真、结果分析和可视化
simulation_chain = SequentialChain(
    chains=[
        TransformChain(
            input_variables=["component_params"],
            output_variables=["component_params_clean"],
            transform=lambda inputs: {"component_params_clean": parse_json_response(inputs["component_params"])}
        ),
        # 3.1 根据component_params生成网表文件
        TransformChain(
            input_variables=["component_params_clean"],
            output_variables=["netlog"],
            transform=lambda inputs: {
                "netlog": generate_bandpass_filter_tool(
                    order=json.loads(inputs["component_params_clean"])["order"],
                    highpass_r=json.loads(inputs["component_params_clean"])["highpass_r"],
                    highpass_c=json.loads(inputs["component_params_clean"])["highpass_c"],
                    lowpass_r=json.loads(inputs["component_params_clean"])["lowpass_r"],
                    lowpass_c=json.loads(inputs["component_params_clean"])["lowpass_c"],
                    filename="BPF.sp"
                )
            }
        ),
        # 3.2 运行HSPICE仿真
        TransformChain(
            input_variables=[],
            output_variables=["simlog"],
            transform=lambda _: {"simlog": run_simulation(" ")}
        ),
        # 3.3 分析仿真结果
        TransformChain(
            input_variables=[],
            output_variables=["analysis"],
            transform=lambda _: {"analysis": analyze_results(" ")}
        ),
        # # 3.4 生成频率响应图
        # TransformChain(
        #     input_variables=[],
        #     output_variables=["plot"],
        #     transform=lambda _: {"plot": visualize(" ")}
        # )
    ],
    # 此处把输入变量修改为component_params，确保上面的transformChain能拿到数据
    input_variables=["component_params"],
    output_variables=["component_params_clean","netlog", "simlog", "analysis"]
)

# --------------------- 完整流程 ---------------------
full_chain = SequentialChain(
    chains=[design_spec_chain, calc_chain, simulation_chain],
    input_variables=["input"],
    output_variables=["design_spec", "component_params", "component_params_clean","netlog", "simlog", "analysis"],
    verbose=True
)

# --------------------- 执行示例 ---------------------
if __name__ == "__main__":
    # 保存设计规范（可以将设计规范保存为JSON文件供后续使用，比如仿真结果验证）
    design_query = """需要设计RC带通滤波器：
    - 中心频率：600Hz
    - 带宽：8000Hz
    """

    result = full_chain.invoke({"input": design_query})

    # 打印各阶段输出
    print("\n设计规范：")
    try:
        print(json.dumps(json.loads(parse_json_response(result["design_spec"])), indent=2, ensure_ascii=False))
    except Exception as e:
        print("设计规范输出有误：", result["design_spec"])

    print("\n元件参数：")
    try:
        print(json.dumps(json.loads(parse_json_response(result["component_params"])), indent=2, ensure_ascii=False))
    except Exception as e:
        print("元件参数输出有误：", result["component_params"])

    print("\n执行日志：")
    print(f"网表生成：{result.get('netlog', '无输出')}")
    print(f"仿真结果：{result.get('simlog', '无输出')}")
    print(f"指标验证：{result.get('analysis', '无输出')}")
    # print(f"可视化：{result.get('plot', '无输出')}")
